from syllabus_prerequisite import construct_node_prerequisite_matrix
from summarization_interface import Material
from collections import namedtuple
import itertools
import re
import numpy as np
from pathlib import Path
from collections import Counter

PreRequisiteData = namedtuple('PreRequisiteData', ['basic_concepts', 'learning_concepts'])

khan_dir = Path(__file__).parent.parent.parent / 'datasets/cosmopedia/khanacademy'

def get_concept_set(file_name):
    material = Material(file_name, khan_dir / 'links' / file_name, khan_dir / 'sentences' / file_name)

    concepts_counter = Counter()
    for sentence in material.sentences:
        for pre, ind, rtype, post, position in sentence.links:
            concepts_counter.update([pre, post])

    return concepts_counter


def get_concept_sets(learning_sequence):
    concepts_sequence = []
    all_counter = Counter()
    for items in learning_sequence:
        concepts_counter = Counter()
        if type(items) is str:
            items = [items]

        for item in items:
            concepts_counter |= get_concept_set(item)
        concepts_sequence.append(concepts_counter)

        all_counter |= concepts_counter
    
    top_concepts = all_counter.most_common(200)
    print(top_concepts)

    filtered_sequence = []
    for concepts in concepts_sequence:
        filtered_sequence.append(Counter({k: v for k, v in concepts.items() if k in top_concepts}))

    return filtered_sequence


# 标注流程，人先检验，然后标注（无pre关系，），然后用大模型进行判断，
def get_prerequisite_matrix(learning_sequence):
    prerequisite_matrix = {}
    pre_concepts = Counter()
    for concepts_counter in get_concept_sets(learning_sequence):
        learned_concepts = concepts_counter & pre_concepts
        learning_concepts = concepts_counter - pre_concepts
        
        for pair in itertools.product(learned_concepts, learning_concepts):
            prerequisite_matrix[pair] = prerequisite_matrix.get(pair, 0) + pre_concepts[pair[0]] + concepts_counter[pair[1]]
        
        pre_concepts.update(concepts_counter)

    for sorted_pair in sorted(prerequisite_matrix, key=lambda x: prerequisite_matrix[x], reverse=True)[:100]:
        print(sorted_pair)

    return prerequisite_matrix


def main():
    expr = experiment_3to9_math()
    # expr = experiment_12to15_math()
    # expr = experiment_reading_vocabulary()
    # expr = experiment_khan_foundation_rsm()

    matrix = get_prerequisite_matrix(expr)


def experiment_3to9_math():
    # 实验1：一年级作为2-9年级的基础（Eureka MathEngageNY
    learning_sequence = [
        '3rd grade (Eureka MathEngageNY)',
        '4th grade (Eureka MathEngageNY)',
        '5th grade (Eureka MathEngageNY)',
        '6th grade (Eureka MathEngageNY)',
        '7th grade (Eureka MathEngageNY)',
        '8th grade (Eureka MathEngageNY)',
        # '9th grade (Eureka MathEngageNY)',
    ]

    concept_groups = []; group = []
    for line in (Path(__file__).parent / 'prerequisite_annotation/experiment_3to9_math').read_text().split('\n'):
        if not line: continue
        if ' - ' not in line:
            if not group: continue
            concept_groups.append(group)
            group = []
            continue

        word, _ = line.split(' - ')
        group.append(word)
    else:
        concept_groups.append(group)

    learning_sequence = [
        Material(file_name, khan_dir / 'links' / file_name, khan_dir /'sentences' / file_name)
        for file_name in learning_sequence
    ]

    return learning_sequence, concept_groups


def experiment_reading_vocabulary():
    # 实验2：一年级作为基础，2-9年级作为进阶 reading & vocabulary
    learning_sequence = [
        '1st grade reading & vocabulary',
        '2nd grade reading & vocabulary',
        '3rd grade reading & vocabulary',
        '4th grade reading & vocabulary',
        '5th grade reading & vocabulary',
        '6th grade reading & vocabulary',
        '7th grade reading & vocabulary',
        '8th grade reading & vocabulary',
        '9th grade reading & vocabulary',
    ]

    return learning_sequence


def experiment_12to15_math():
    # 实验3：Middle school作为基础，支持high school (medicine) ...
    learning_sequence = [
        ('Middle school biology', 'Middle school physics'),
        ('High school biology', 'High school physics'),
    ]

    return learning_sequence


def experiment_khan_foundation_rsm():
    # 实验4：所有khan -> foundation -> rsm的抽取
    learning_sequence = [
        'khan',
        'foundation of database',
        'resource space model',
    ]
    return learning_sequence


def experiment_university_course():
    '''
    返回一个矩阵
    '''
    base_dir = Path(__file__).parent.parent.parent / 'datasets/prerequisite/benchmarks/benchmarks/conventional_setting/UniversityCourse'

    concept_to_id = {}; id_to_concept = {}
    for line in (base_dir / 'concepts.tsv').read_text(encoding='utf-8').split('\n'):
        if not line: continue
        concept, concept_id = line.split('\t')
        concept = re.sub('\(.*\)', '', concept).strip()

        concept_to_id[concept] = int(concept_id)
        id_to_concept[int(concept_id)] = concept

    pattern = '(' + '|'.join(sorted(concept_to_id.keys(), key=lambda x: -len(x))) + ')'
    
    subject_prereq = {}
    constraint_prereq = {}
    co_occurrence_prereq = {}
    prereq_matrix = np.zeros((int(concept_id) + 1, int(concept_id) + 1))

    for line in (base_dir / 'descriptions.tsv').read_text(encoding='utf-8').split('\n'):
        if not line: continue

        concept, description = line.split('\t')
        description = description.lower()
        concept = re.sub('\(.*\)', '', concept).strip()
        # 出现在description中的concept都是他的先序
        for matched_concept in re.findall(pattern, description)[:4]:
            if matched_concept == concept: continue
            if matched_concept not in concept_to_id: continue

            if description.startswith(concept) or description.startswith('a'):
                subject_prereq[(matched_concept, concept)] = 1
            else:
                constraint_prereq[(matched_concept, concept)] = 1

        for matched_concept in re.findall(pattern, description):
            if matched_concept == concept: continue
            if matched_concept not in concept_to_id: continue

            co_occurrence_prereq[(matched_concept, concept)] = 1
            prereq_matrix[concept_to_id[matched_concept], concept_to_id[concept]] = 1

    reference_1 = prereq_matrix + np.matmul(prereq_matrix, prereq_matrix)
    reference_1_keys = set()
    for i in range(prereq_matrix.shape[0]):
        for j in range(prereq_matrix.shape[0]):
            if reference_1[i, j]:
                reference_1_keys.add((id_to_concept[i], id_to_concept[j]))

    reference_2 = reference_1 + np.matmul(reference_1, reference_1)
    reference_2_keys = set()
    for i in range(prereq_matrix.shape[0]):
        for j in range(prereq_matrix.shape[0]):
            if reference_2[i, j]:
                reference_2_keys.add((id_to_concept[i], id_to_concept[j]))

    reference_3 = reference_2 + np.matmul(reference_2, reference_2)
    reference_3_keys = set()
    for i in range(prereq_matrix.shape[0]):
        for j in range(prereq_matrix.shape[0]):
            if reference_3[i, j]:
                reference_3_keys.add((id_to_concept[i], id_to_concept[j]))

    prereq_dict = {}
    for line in (base_dir / 'prerequisites.tsv').read_text().split('\n'):
        if not line: continue

        pre_concept, _, concept = line.split('\t')
        pre_concept = re.sub('\(.*\)', '', pre_concept).strip()
        concept = re.sub('\(.*\)', '', concept).strip()

        prereq_dict[(pre_concept, concept)] = 1
    
    def p_r(name, keys, truth_keys):
        tp = keys & truth_keys
        print(f'{name} Precision: ', round(len(tp) / len(keys), 4), f'  {len(tp)} / {len(keys)}')
        print(f'{name} Recall: ', round(len(tp) / len(truth_keys), 4), f'  {len(tp)} / {len(truth_keys)}')
        print()

    p_r('co_occurrence', co_occurrence_prereq.keys(), prereq_dict.keys())
    p_r('reference_1', reference_1_keys, prereq_dict.keys())
    p_r('reference_2', reference_2_keys, prereq_dict.keys())
    p_r('reference_3', reference_3_keys, prereq_dict.keys())

    # print(prereq_dict.keys() - semantic_prereq.keys())
    p_r('subject', subject_prereq.keys(), co_occurrence_prereq.keys() & prereq_dict.keys())
    p_r('constraint', constraint_prereq.keys(), co_occurrence_prereq.keys() & prereq_dict.keys())



if __name__ == '__main__':
    # main()
    experiment_university_course()